package com.yanyeori.framework.websocket.server;

import com.yanyeori.framework.core.util.JacksonUtil;
import com.yanyeori.framework.websocket.message.WsMessage;
import lombok.extern.slf4j.Slf4j;

import javax.websocket.Session;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * WebSocket 工具
 *
 * @author chenkuan
 * @since 2021/6/3
 */
@Slf4j
public class WebSocketUtil {

    private static final Map<String, Session> SESSION_MAP = new ConcurrentHashMap<>();
    private static final Map<String, List<String>> SUBSCRIBE_ID_TOPIC_MAP = new ConcurrentHashMap<>();
    private static final Map<String, List<String>> SUBSCRIBE_TOPIC_ID_MAP = new ConcurrentHashMap<>();

    /**
     * 新的session连接
     */
    public static void addSession(Session session) {
        SESSION_MAP.put(session.getId(), session);
    }

    /**
     * session断开连接，移除session及session订阅的所有主题
     */
    public static void removeSession(Session session) {
        String sessionId = session.getId();
        SESSION_MAP.remove(sessionId);
        for (String topic : SUBSCRIBE_ID_TOPIC_MAP.getOrDefault(sessionId, new ArrayList<>())) {
            SUBSCRIBE_TOPIC_ID_MAP.get(topic).removeIf(sessionId::equals);
        }
        SUBSCRIBE_ID_TOPIC_MAP.remove(sessionId);
    }

    /**
     * 获取当前连接session数
     */
    public static int sessionSize() {
        return SESSION_MAP.size();
    }

    /**
     * 发送同步消息
     *
     * @param message 消息内容
     * @param topics  主题
     */
    public static void send(Object message, String... topics) throws IOException {
        send(message, false, topics);
    }

    /**
     * 发送异步消息
     *
     * @param message 消息内容
     * @param topics  主题
     */
    public static void sendAsync(Object message, String... topics) throws IOException {
        send(message, true, topics);
    }

    /**
     * 发送消息
     *
     * @param message 消息内容
     * @param isAsync 是否异步
     * @param topics  主题
     */
    private static void send(Object message, boolean isAsync, String... topics) throws IOException {
        for (String topic : topics) {
            List<String> sessionIds = SUBSCRIBE_TOPIC_ID_MAP.get(topic);
            if (sessionIds == null || sessionIds.isEmpty()) {
                log.debug("No client subscribed to topic [{}], the message was not sent", topic);
                continue;
            }
            String text = JacksonUtil.toJson(new WsMessage(topic, message));
            for (String sessionId : sessionIds) {
                Session session = SESSION_MAP.get(sessionId);
                if (session == null) {
                    log.debug("Client sessionId [{}] is not online and the message was not sent", sessionId);
                    continue;
                }
                if (isAsync) {
                    session.getAsyncRemote().sendText(text);
                } else {
                    session.getBasicRemote().sendText(text);
                }
                log.debug("Send message: sessionId=[{}], message=[{}]", sessionId, text);
            }
        }
    }

    /**
     * 订阅主题
     */
    public static String subscribe(String sessionId, String topic) {
        if (SUBSCRIBE_ID_TOPIC_MAP.get(sessionId) != null && SUBSCRIBE_ID_TOPIC_MAP.get(sessionId).contains(topic)) {
            return String.format("Topic [%s] has been subscribed. Please do not subscribe again", topic);
        }
        SUBSCRIBE_ID_TOPIC_MAP.computeIfAbsent(sessionId, k -> new ArrayList<>()).add(topic);
        SUBSCRIBE_TOPIC_ID_MAP.computeIfAbsent(topic, k -> new ArrayList<>()).add(sessionId);
        return String.format("Successfully subscribed to topic: %s", topic);
    }

    /**
     * 取消订阅主题
     */
    public static String unsubscribe(String sessionId, String topic) {
        if (SUBSCRIBE_ID_TOPIC_MAP.get(sessionId) == null || !SUBSCRIBE_ID_TOPIC_MAP.get(sessionId).contains(topic)) {
            return String.format("The topic [%s] is not subscribed. Please subscribe to the topic first", topic);
        }
        SUBSCRIBE_ID_TOPIC_MAP.getOrDefault(sessionId, new ArrayList<>()).removeIf(topic::equals);
        SUBSCRIBE_TOPIC_ID_MAP.getOrDefault(topic, new ArrayList<>()).removeIf(sessionId::equals);
        return String.format("Unsubscribed topic: %s", topic);
    }
}
